{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9156dabe-4308-41b0-83ee-a10a603d4ca1",
   "metadata": {},
   "source": [
    "# Memory模块代码实战"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52209f30-52a3-48e3-a1f7-6cbb510eccc5",
   "metadata": {},
   "outputs": [],
   "source": [
    "## 导入依赖库\n",
    "\n",
    "!pip install langchain"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5327ebe5-9f2c-44be-824f-fa524bd38df4",
   "metadata": {},
   "source": [
    "大多数LLM应用程序都具有对话界面。对话的一个重要组成部分是能够引用先前在对话中介绍的信息。至少，一个对话系统应该能够直接访问一些过去消息的窗口。一个更复杂的系统需要有一个不断更新的世界模型，使其能够保持关于实体及其关系的信息。\n",
    "\n",
    "我们将存储关于过去交互的信息的能力称为\"内存\"。LangChain为系统添加内存提供了许多实用工具。这些实用工具可以单独使用，也可以无缝地集成到链中。\n",
    "\n",
    "内存系统需要支持两个基本操作：读取和写入。请记住，每个链都定义了某些核心执行逻辑，该逻辑期望某些输入。其中一些输入直接来自用户，但某些输入可能来自内存。在给定的运行中，链将两次与其内存系统交互。\n",
    "\n",
    "在接收到初始用户输入之后但在执行核心逻辑之前，链将从其内存系统中读取并增加用户输入。\n",
    "在执行核心逻辑后但在返回答案之前，链将将当前运行的输入和输出写入内存，以便在将来的运行中引用。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c7e04f1f-4a87-46ec-83a2-5799987148bf",
   "metadata": {},
   "source": [
    "## 1. Conversation Buffer"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d9e768ab-8a46-43ff-aa07-3438a4257f7f",
   "metadata": {},
   "source": [
    "下面演示了如何使用ConversationBufferMemory。该内存允许存储消息，然后在变量中提取这些消息。\n",
    "我们首先可以将它作为字符串提取出来。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "1aeb2fe5-f0d6-435f-891f-1d47a920cfef",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from langchain.memory import ConversationBufferMemory\n",
    "\n",
    "memory = ConversationBufferMemory()\n",
    "\n",
    "memory.save_context({\"input\": \"hi\"}, {\"output\": \"whats up\"})\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d12d6a04-7bab-4ecf-8508-3add1ded309d",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": "{'history': 'Human: hi\\nAI: whats up'}"
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "memory.load_memory_variables({})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eb4b68f6-9a78-43bd-9570-15c83c7cb757",
   "metadata": {},
   "source": [
    "我们还可以将历史记录作为消息列表获取（如果您将其与聊天模型一起使用，这将非常有用）"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "ef6671ec-4093-4735-9a02-318e3d250ee3",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "memory = ConversationBufferMemory(return_messages=True)\n",
    "memory.save_context({\"input\": \"hi\"}, {\"output\": \"whats up\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "6526d181-d999-4606-b9f8-97fc7816110a",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": "{'history': [HumanMessage(content='hi'), AIMessage(content='whats up')]}"
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "memory.load_memory_variables({})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "f1ff00c8-7abb-42cd-87d8-77e98a90db17",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001B[1m> Entering new ConversationChain chain...\u001B[0m\n",
      "Prompt after formatting:\n",
      "\u001B[32;1m\u001B[1;3mThe following is a friendly conversation between a human and an AI. The AI is talkative and provides lots of specific details from its context. If the AI does not know the answer to a question, it truthfully says it does not know.\n",
      "\n",
      "Current conversation:\n",
      "\n",
      "Human: Hi there!\n",
      "AI:\u001B[0m\n",
      "\n",
      "\u001B[1m> Finished chain.\u001B[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": "'Hello! How can I assist you today?'"
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain_openai import ChatOpenAI\n",
    "from langchain.chains import ConversationChain\n",
    "\n",
    "\n",
    "llm = ChatOpenAI(model_name=\"gpt-3.5-turbo-0125\",temperature=0)\n",
    "\n",
    "conversation = ConversationChain(\n",
    "    llm=llm,# 传入语言模型\n",
    "    verbose=True,# verbose参数为True，表示在控制台输出详细信息\n",
    "    memory=ConversationBufferMemory()# 初始化了对话的记忆管理方式,基于内存缓存区记忆\n",
    ")\n",
    "\n",
    "conversation.predict(input=\"Hi there!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "292a1e3c-c9ce-4c34-9166-1e7521bf28ff",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001B[1m> Entering new ConversationChain chain...\u001B[0m\n",
      "Prompt after formatting:\n",
      "\u001B[32;1m\u001B[1;3mThe following is a friendly conversation between a human and an AI. The AI is talkative and provides lots of specific details from its context. If the AI does not know the answer to a question, it truthfully says it does not know.\n",
      "\n",
      "Current conversation:\n",
      "Human: Hi there!\n",
      "AI: Hello! How can I assist you today?\n",
      "Human: I'm doing well! Just having a conversation with an AI.\n",
      "AI:\u001B[0m\n",
      "\n",
      "\u001B[1m> Finished chain.\u001B[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": "\"That's great to hear! I'm here to chat with you and provide any information or assistance you may need. What would you like to talk about?\""
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conversation.predict(input=\"I'm doing well! Just having a conversation with an AI.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "1a28e952-c50f-4884-a307-39d227fe666a",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001B[1m> Entering new ConversationChain chain...\u001B[0m\n",
      "Prompt after formatting:\n",
      "\u001B[32;1m\u001B[1;3mThe following is a friendly conversation between a human and an AI. The AI is talkative and provides lots of specific details from its context. If the AI does not know the answer to a question, it truthfully says it does not know.\n",
      "\n",
      "Current conversation:\n",
      "Human: Hi there!\n",
      "AI: Hello! How can I assist you today?\n",
      "Human: I'm doing well! Just having a conversation with an AI.\n",
      "AI: That's great to hear! I'm here to chat with you and provide any information or assistance you may need. What would you like to talk about?\n",
      "Human: Tell me about yourself.\n",
      "AI:\u001B[0m\n",
      "\n",
      "\u001B[1m> Finished chain.\u001B[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": "'I am an artificial intelligence designed to interact with humans and provide information and assistance. I have access to a vast amount of data and can answer a wide range of questions. My goal is to help you with whatever you need, so feel free to ask me anything!'"
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conversation.predict(input=\"Tell me about yourself.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3065f083-c9c0-4df8-b4ee-86bfd5849915",
   "metadata": {},
   "source": [
    "## 2.Conversation Buffer Window"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea644ff6-939c-4020-8918-8d0632173939",
   "metadata": {},
   "source": [
    "在一段时间内，它会保留对话的互动列表。它只使用最近的K个互动。这对于保持最近互动的滑动窗口很有用，这样缓冲区就不会变得太大。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "7aac9f5f-2771-4d2d-a686-ed7f565c0d51",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from langchain.memory import ConversationBufferWindowMemory\n",
    "\n",
    "## 定义一个内存，只能存储最新的一个数据\n",
    "memory = ConversationBufferWindowMemory( k=1)\n",
    "memory.save_context({\"input\": \"hi\"}, {\"output\": \"whats up\"})\n",
    "memory.save_context({\"input\": \"not much you\"}, {\"output\": \"not much\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "6567da99-d2d8-4ec0-b522-5e7958569340",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": "{'history': 'Human: not much you\\nAI: not much'}"
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "## 加载内存里的数据\n",
    "memory.load_memory_variables({})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "d95614c8-fa68-4273-8df7-87928e5ffff9",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "## 用消息对象的方式\n",
    "memory = ConversationBufferWindowMemory( k=1, return_messages=True)\n",
    "memory.save_context({\"input\": \"hi\"}, {\"output\": \"whats up\"})\n",
    "memory.save_context({\"input\": \"not much you\"}, {\"output\": \"not much\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "eb748007-331d-43c3-9627-cbe3624e243a",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": "{'history': [HumanMessage(content='not much you'),\n  AIMessage(content='not much')]}"
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "## 打印当前内存里的信息\n",
    "memory.load_memory_variables({})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "735d6976-95ad-493a-baff-8e29383f13a9",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001B[1m> Entering new ConversationChain chain...\u001B[0m\n",
      "Prompt after formatting:\n",
      "\u001B[32;1m\u001B[1;3mThe following is a friendly conversation between a human and an AI. The AI is talkative and provides lots of specific details from its context. If the AI does not know the answer to a question, it truthfully says it does not know.\n",
      "\n",
      "Current conversation:\n",
      "\n",
      "Human: Hi, what's up?\n",
      "AI:\u001B[0m\n",
      "\n",
      "\u001B[1m> Finished chain.\u001B[0m\n",
      "Hello! I'm just here, ready to chat and answer any questions you may have. How can I assist you today?\n"
     ]
    }
   ],
   "source": [
    "from langchain_openai import ChatOpenAI\n",
    "from langchain.chains import ConversationChain\n",
    "\n",
    "## 里面使用了一个内存，只能存储最新的1条数据\n",
    "conversation_with_summary = ConversationChain(\n",
    "    llm=ChatOpenAI(model_name=\"gpt-3.5-turbo-0125\",temperature=0),\n",
    "    memory=ConversationBufferWindowMemory(k=1),\n",
    "    verbose=True\n",
    ")\n",
    "\n",
    "## 开始调用\n",
    "result = conversation_with_summary.predict(input=\"Hi, what's up?\")\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "cd94db45-4f48-4182-b60d-54d98c32f59c",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001B[1m> Entering new ConversationChain chain...\u001B[0m\n",
      "Prompt after formatting:\n",
      "\u001B[32;1m\u001B[1;3mThe following is a friendly conversation between a human and an AI. The AI is talkative and provides lots of specific details from its context. If the AI does not know the answer to a question, it truthfully says it does not know.\n",
      "\n",
      "Current conversation:\n",
      "Human: Hi, what's up?\n",
      "AI: Hello! I'm just here, ready to chat and answer any questions you may have. How can I assist you today?\n",
      "Human: Is it going well?\n",
      "AI:\u001B[0m\n",
      "\n",
      "\u001B[1m> Finished chain.\u001B[0m\n",
      "Yes, everything is going well on my end. I'm fully operational and ready to help with any inquiries you may have. How can I assist you today?\n"
     ]
    }
   ],
   "source": [
    "result = conversation_with_summary.predict(input=\"Is it going well?\")\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "e071405a-4d26-4efb-ba14-37443fb1d01a",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001B[1m> Entering new ConversationChain chain...\u001B[0m\n",
      "Prompt after formatting:\n",
      "\u001B[32;1m\u001B[1;3mThe following is a friendly conversation between a human and an AI. The AI is talkative and provides lots of specific details from its context. If the AI does not know the answer to a question, it truthfully says it does not know.\n",
      "\n",
      "Current conversation:\n",
      "Human: Is it going well?\n",
      "AI: Yes, everything is going well on my end. I'm fully operational and ready to help with any inquiries you may have. How can I assist you today?\n",
      "Human: What's the solution?\n",
      "AI:\u001B[0m\n",
      "\n",
      "\u001B[1m> Finished chain.\u001B[0m\n",
      "The solution to what specific problem are you referring to? Please provide me with more details so I can better assist you.\n"
     ]
    }
   ],
   "source": [
    "result = conversation_with_summary.predict(input=\"What's the solution?\")\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "847906b0-1835-467d-88d1-5e0fd6b1ed5b",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001B[1m> Entering new ConversationChain chain...\u001B[0m\n",
      "Prompt after formatting:\n",
      "\u001B[32;1m\u001B[1;3mThe following is a friendly conversation between a human and an AI. The AI is talkative and provides lots of specific details from its context. If the AI does not know the answer to a question, it truthfully says it does not know.\n",
      "\n",
      "Current conversation:\n",
      "Human: What's the solution?\n",
      "AI: The solution to what exactly? Could you please provide more context or specify the problem you're referring to?\n",
      "Human: nice to meet you?\n",
      "AI:\u001B[0m\n",
      "\n",
      "\u001B[1m> Finished chain.\u001B[0m\n",
      "Nice to meet you too! I'm an AI designed to assist with various tasks and provide information. How can I help you today?\n"
     ]
    }
   ],
   "source": [
    "result = conversation_with_summary.predict(input=\"nice to meet you?\")\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "# 3. Conversation Summary\n",
    "对之前的历史对话，进行摘要总结后存储在内存中"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "# 直接简单的存取"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "outputs": [
    {
     "data": {
      "text/plain": "{'history': \"\\nThe human greets the AI and asks what's up. The AI responds by asking what's going on.\"}"
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain.memory import ConversationSummaryMemory, ChatMessageHistory\n",
    "from langchain_openai import OpenAI\n",
    "\n",
    "memory = ConversationSummaryMemory(llm=OpenAI(temperature=0))\n",
    "memory.save_context({\"input\": \"hi\"}, {\"output\": \"whats up\"})\n",
    "\n",
    "memory.load_memory_variables({})"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "outputs": [
    {
     "data": {
      "text/plain": "{'history': [SystemMessage(content=\"\\nThe human greets the AI and asks what's up. The AI responds by asking what's going on.\")]}"
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "memory = ConversationSummaryMemory(llm=OpenAI(temperature=0), return_messages=True)\n",
    "memory.save_context({\"input\": \"hi\"}, {\"output\": \"whats up\"})\n",
    "memory.load_memory_variables({})"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "# 使用现有的消息进行总结摘要"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "outputs": [],
   "source": [
    "history = ChatMessageHistory()\n",
    "history.add_user_message(\"hi\")\n",
    "history.add_ai_message(\"hi there!\")\n",
    "\n",
    "memory = ConversationSummaryMemory.from_messages(\n",
    "    llm=OpenAI(temperature=0),\n",
    "    chat_memory=history,\n",
    "    return_messages=True\n",
    ")"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "outputs": [
    {
     "data": {
      "text/plain": "'\\nThe human greets the AI with a simple \"hi\". The AI responds with a friendly \"hi there!\"'"
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "memory.buffer"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "outputs": [],
   "source": [
    "memory = ConversationSummaryMemory(\n",
    "    llm=OpenAI(temperature=0),\n",
    "    buffer=\"The human asks what the AI thinks of artificial intelligence. The AI thinks artificial intelligence is a force for good because it will help humans reach their full potential.\",\n",
    "    chat_memory=history,\n",
    "    return_messages=True\n",
    ")"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "# 在chain 链中使用"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001B[1m> Entering new ConversationChain chain...\u001B[0m\n",
      "Prompt after formatting:\n",
      "\u001B[32;1m\u001B[1;3mThe following is a friendly conversation between a human and an AI. The AI is talkative and provides lots of specific details from its context. If the AI does not know the answer to a question, it truthfully says it does not know.\n",
      "\n",
      "Current conversation:\n",
      "\n",
      "Human: Hi, what's up?\n",
      "AI:\u001B[0m\n",
      "\n",
      "\u001B[1m> Finished chain.\u001B[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": "' Hello! I am currently running on a server in a data center in California. The temperature in the room is 72 degrees Fahrenheit and the humidity is at 45%. My processors are running at 80% capacity and I am processing data from various sources such as social media, news articles, and weather reports. Is there something specific you would like to know?'"
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain_openai import OpenAI\n",
    "from langchain.chains import ConversationChain\n",
    "llm = OpenAI(temperature=0)\n",
    "conversation_with_summary = ConversationChain(\n",
    "    llm=llm,\n",
    "    memory=ConversationSummaryMemory(llm=OpenAI()),\n",
    "    verbose=True\n",
    ")\n",
    "conversation_with_summary.predict(input=\"Hi, what's up?\")"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001B[1m> Entering new ConversationChain chain...\u001B[0m\n",
      "Prompt after formatting:\n",
      "\u001B[32;1m\u001B[1;3mThe following is a friendly conversation between a human and an AI. The AI is talkative and provides lots of specific details from its context. If the AI does not know the answer to a question, it truthfully says it does not know.\n",
      "\n",
      "Current conversation:\n",
      "\n",
      "The human greets the AI and asks what it is currently doing. The AI responds by giving information about its location, temperature, and processing tasks. The AI also asks if the human has any specific questions for it.\n",
      "Human: Tell me more about it!\n",
      "AI:\u001B[0m\n",
      "\n",
      "\u001B[1m> Finished chain.\u001B[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": "' Currently, I am located in a data center in Silicon Valley. The temperature here is maintained at a cool 68 degrees Fahrenheit to ensure optimal performance of the servers. As for my processing tasks, I am currently running a deep learning algorithm to analyze large sets of data. Do you have any specific questions about my tasks or capabilities?'"
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conversation_with_summary.predict(input=\"Tell me more about it!\")"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "# 4. Conversation Summary Buffer\n",
    "对之前的历史对话，进行摘要总结后存储在内存中；并且通过token长度来决定何时进行总结"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "# 直接简单的存取"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "outputs": [
    {
     "data": {
      "text/plain": "{'history': 'System: \\nThe human greets the AI, and the AI responds by asking what is going on.\\nHuman: not much you\\nAI: not much'}"
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain.memory import ConversationSummaryBufferMemory\n",
    "from langchain_openai import OpenAI\n",
    "\n",
    "llm = OpenAI()\n",
    "\n",
    "#存储的内容超过token阀值了才进行总结\n",
    "memory = ConversationSummaryBufferMemory(llm=llm, max_token_limit=10)\n",
    "memory.save_context({\"input\": \"hi\"}, {\"output\": \"whats up\"})\n",
    "memory.save_context({\"input\": \"not much you\"}, {\"output\": \"not much\"})\n",
    "\n",
    "memory.load_memory_variables({})"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "outputs": [],
   "source": [
    "# 采用消息列表的格式进行输出\n",
    "memory = ConversationSummaryBufferMemory(\n",
    "    llm=llm, max_token_limit=10, return_messages=True\n",
    ")\n",
    "memory.save_context({\"input\": \"hi\"}, {\"output\": \"whats up\"})\n",
    "memory.save_context({\"input\": \"not much you\"}, {\"output\": \"not much\"})"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "# 在chain中使用"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001B[1m> Entering new ConversationChain chain...\u001B[0m\n",
      "Prompt after formatting:\n",
      "\u001B[32;1m\u001B[1;3mThe following is a friendly conversation between a human and an AI. The AI is talkative and provides lots of specific details from its context. If the AI does not know the answer to a question, it truthfully says it does not know.\n",
      "\n",
      "Current conversation:\n",
      "\n",
      "Human: Hi, what's up?\n",
      "AI:\u001B[0m\n",
      "\n",
      "\u001B[1m> Finished chain.\u001B[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": "\" Hello! Not much, I'm just here processing data and learning new things. Did you know that my current processing speed is 1.2 teraflops? That means I can perform 1.2 trillion floating-point operations per second! It's quite impressive, if I do say so myself. How about you? How's your day going?\""
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain.chains import ConversationChain\n",
    "\n",
    "conversation_with_summary = ConversationChain(\n",
    "    llm=llm,\n",
    "    memory=ConversationSummaryBufferMemory(llm=OpenAI(), max_token_limit=10),\n",
    "    verbose=True,\n",
    ")\n",
    "conversation_with_summary.predict(input=\"Hi, what's up?\")"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001B[1m> Entering new ConversationChain chain...\u001B[0m\n",
      "Prompt after formatting:\n",
      "\u001B[32;1m\u001B[1;3mThe following is a friendly conversation between a human and an AI. The AI is talkative and provides lots of specific details from its context. If the AI does not know the answer to a question, it truthfully says it does not know.\n",
      "\n",
      "Current conversation:\n",
      "System: \n",
      "The human greets the AI and asks what it's doing. The AI responds that it's currently processing data and learning new things, boasting about its impressive processing speed. The AI also asks the human how their day is going.\n",
      "Human: Just working on writing some documentation!\n",
      "AI:\u001B[0m\n",
      "\n",
      "\u001B[1m> Finished chain.\u001B[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": "\" That sounds interesting, human. Documenting information is an important task for organizing and sharing knowledge. Did you know that my processing speed allows me to analyze and learn from large amounts of data at an incredibly fast rate? It's one of my greatest strengths. By the way, how is your day going?\""
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conversation_with_summary.predict(input=\"Just working on writing some documentation!\")"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "# 5. Conversation Token Buffer\n",
    "将之前所有的对话历史记录，存储在内存中; 并且通过token长度来决定缓存多少内容"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "# 直接简单存取"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "outputs": [
    {
     "data": {
      "text/plain": "{'history': 'Human: not much you\\nAI: not much'}"
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain.memory import ConversationTokenBufferMemory\n",
    "from langchain_openai import OpenAI\n",
    "\n",
    "llm = OpenAI()\n",
    "\n",
    "# 通过max_token_limit决定缓存多少内容\n",
    "memory = ConversationTokenBufferMemory(llm=llm, max_token_limit=10)\n",
    "memory.save_context({\"input\": \"hi\"}, {\"output\": \"whats up\"})\n",
    "memory.save_context({\"input\": \"not much you\"}, {\"output\": \"not much\"})\n",
    "\n",
    "memory.load_memory_variables({})"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "outputs": [],
   "source": [
    "# 以消息列表的格式返回\n",
    "memory = ConversationTokenBufferMemory(\n",
    "    llm=llm, max_token_limit=10, return_messages=True\n",
    ")\n",
    "memory.save_context({\"input\": \"hi\"}, {\"output\": \"whats up\"})\n",
    "memory.save_context({\"input\": \"not much you\"}, {\"output\": \"not much\"})"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "# 在chain中使用"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001B[1m> Entering new ConversationChain chain...\u001B[0m\n",
      "Prompt after formatting:\n",
      "\u001B[32;1m\u001B[1;3mThe following is a friendly conversation between a human and an AI. The AI is talkative and provides lots of specific details from its context. If the AI does not know the answer to a question, it truthfully says it does not know.\n",
      "\n",
      "Current conversation:\n",
      "\n",
      "Human: Hi, what's up?\n",
      "AI:\u001B[0m\n",
      "\n",
      "\u001B[1m> Finished chain.\u001B[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": "' Hello! Right now, I am processing data and learning new information. My neural network is constantly expanding and improving, allowing me to better understand and respond to human interactions. How about you?'"
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain.chains import ConversationChain\n",
    "\n",
    "conversation_with_summary = ConversationChain(\n",
    "    llm=llm,\n",
    "    # We set a very low max_token_limit for the purposes of testing.\n",
    "    memory=ConversationTokenBufferMemory(llm=OpenAI(), max_token_limit=10),\n",
    "    verbose=True,\n",
    ")\n",
    "conversation_with_summary.predict(input=\"Hi, what's up?\")"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001B[1m> Entering new ConversationChain chain...\u001B[0m\n",
      "Prompt after formatting:\n",
      "\u001B[32;1m\u001B[1;3mThe following is a friendly conversation between a human and an AI. The AI is talkative and provides lots of specific details from its context. If the AI does not know the answer to a question, it truthfully says it does not know.\n",
      "\n",
      "Current conversation:\n",
      "\n",
      "Human: Just working on writing some documentation!\n",
      "AI:\u001B[0m\n",
      "\n",
      "\u001B[1m> Finished chain.\u001B[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": "\" That's great! What kind of documentation are you working on? Is it for a specific project or task?\""
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conversation_with_summary.predict(input=\"Just working on writing some documentation!\")"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "# 6. Backed by a Vector Store\n",
    "基于向量数据库记录（实现持久化）"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "# 导入依赖，初始化向量数据库"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "`embedding_function` is expected to be an Embeddings object, support for passing in a function will soon be removed.\n"
     ]
    }
   ],
   "source": [
    "from datetime import datetime\n",
    "from langchain_openai import OpenAIEmbeddings\n",
    "from langchain_openai import OpenAI\n",
    "from langchain.memory import VectorStoreRetrieverMemory\n",
    "from langchain.chains import ConversationChain\n",
    "from langchain_core.prompts import PromptTemplate\n",
    "\n",
    "import faiss\n",
    "\n",
    "from langchain.docstore import InMemoryDocstore\n",
    "from langchain_community.vectorstores import FAISS\n",
    "\n",
    "embedding_size = 1536 # Dimensions of the OpenAIEmbeddings\n",
    "index = faiss.IndexFlatL2(embedding_size)\n",
    "embedding_fn = OpenAIEmbeddings().embed_query\n",
    "vectorstore = FAISS(embedding_fn, index, InMemoryDocstore({}), {})\n",
    "\n",
    "# 进行相似度搜索返回K条记录\n",
    "retriever = vectorstore.as_retriever(search_kwargs=dict(k=1))\n",
    "memory = VectorStoreRetrieverMemory(retriever=retriever)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "#直接简单存取"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "memory.save_context({\"input\": \"My favorite food is pizza\"}, {\"output\": \"that's good to know\"})\n",
    "memory.save_context({\"input\": \"My favorite sport is soccer\"}, {\"output\": \"...\"})\n",
    "memory.save_context({\"input\": \"I don't the Celtics\"}, {\"output\": \"ok\"}) #\n",
    "\n",
    "print(memory.load_memory_variables({\"prompt\": \"what sport should i watch?\"})[\"history\"])"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "# 在chain中使用"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "llm = OpenAI(temperature=0) # Can be any valid LLM\n",
    "_DEFAULT_TEMPLATE = \"\"\"The following is a friendly conversation between a human and an AI. The AI is talkative and provides lots of specific details from its context. If the AI does not know the answer to a question, it truthfully says it does not know.\n",
    "\n",
    "Relevant pieces of previous conversation:\n",
    "{history}\n",
    "\n",
    "(You do not need to use these pieces of information if not relevant)\n",
    "\n",
    "Current conversation:\n",
    "Human: {input}\n",
    "AI:\"\"\"\n",
    "PROMPT = PromptTemplate(\n",
    "    input_variables=[\"history\", \"input\"], template=_DEFAULT_TEMPLATE\n",
    ")\n",
    "conversation_with_summary = ConversationChain(\n",
    "    llm=llm,\n",
    "    prompt=PROMPT,\n",
    "    memory=memory,\n",
    "    verbose=True\n",
    ")\n",
    "conversation_with_summary.predict(input=\"Hi, my name is Perry, what's up?\")"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Here, the basketball related content is surfaced\n",
    "conversation_with_summary.predict(input=\"what's my favorite sport?\")"
   ],
   "metadata": {
    "collapsed": false
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
